#!/bin/bash

# 安装Nginx
sudo apt update
sudo apt install -y nginx

# 定义环境变量名称
software_dir="stable-diffusion"
venv_name="amd${software_dir}"


# 显卡型号与HSA_OVERRIDE_GFX_VERSION的对应关系
declare -A version_map=(
    ["1"]="RX 7900 XTX, RX 7900 XT, RX 7900 GRE, PRO W7900, PRO W7800"
    ["2"]="RX 7800 XT, RX 7700 XT, PRO W7700"
    ["3"]="RX 7600"
    ["4"]="RX 6900 XT, RX 6950 XT, RX 6800 XT, RX 6800, PRO V620, PRO W6800"
    ["5"]="RX 6700 XT, RX 6750 XT, RX 6700"
    ["6"]="RX 6600 XT, RX 6600, RX 6650 XT, PRO W6600"
    ["7"]="Pro W5500"
    ["8"]="VII, Pro VII"
)

# HSA_OVERRIDE_GFX_VERSION对应关系
declare -A gfx_version_map=(
    ["1"]="11.0.0"
    ["2"]="11.0.1"
    ["3"]="11.0.2"
    ["4"]="10.3.0"
    ["5"]="10.3.1"
    ["6"]="10.3.2"
    ["7"]="10.1.2"
    ["8"]="9.0.6"
)

# 显示可用的HSA_OVERRIDE_GFX_VERSION选项
echo "请选择对应的显卡版本编号："
for i in {1..8}; do
    echo "$i: ${version_map[$i]}"
done

# 读取用户输入
read -p "输入显卡版本编号，如果显卡不在上面的列表中可以输入相近的型号的数字（或输入 'N' 退出）: " version_choice

# 验证用户输入并查找对应的版本号
if [[ "$version_choice" == "N" || "$version_choice" == "n" ]]; then
    echo "退出脚本。"
    exit 0
elif [[ -v version_map[$version_choice] ]]; then
    found_version=${gfx_version_map[$version_choice]}
    echo "选择的HSA_OVERRIDE_GFX_VERSION: $found_version"
else
    echo "无效的显卡版本编号。"
    echo "请重新运行脚本，并输入有效的编号。"
    exit 1
fi

# 获取本机上网IP地址
IP_ADDR=$(hostname -I | awk '{print $1}')

# 查询系统中的显卡数量
GPU_COUNT=$(lspci | grep -i 'vga compatible controller: Advanced Micro Devices' | wc -l)

# 确认获取到的GPU数量
if [ "$GPU_COUNT" -eq 0 ]; then
    echo "未检测到AMD显卡，脚本退出。"
    exit 1
fi

echo "检测到的AMD显卡数量: $GPU_COUNT"
echo "本机上网IP地址: $IP_ADDR"

sudo rm -rf /etc/nginx/nginx.conf

# 修改Nginx配置文件
sudo bash -c "cat > /etc/nginx/nginx.conf" <<EOL
events {
    worker_connections 1024;
}

http {
    include       mime.types;
    default_type  application/octet-stream;
    
    map \$http_upgrade \$connection_upgrade {
        default upgrade;
        ''      close;
    }

    sendfile        on;
    keepalive_timeout  65;

    upstream sd {
        hash \$remote_addr;
EOL

# 添加每个实例的后端服务器配置
for i in $(seq 0 $((GPU_COUNT - 1))); do
    port=$((7860 + i))
    sudo bash -c "echo '        server $IP_ADDR:$port;' >> /etc/nginx/nginx.conf"
done

sudo bash -c "cat >> /etc/nginx/nginx.conf" <<EOL
    }

    server {
        listen 8080;
        server_name $IP_ADDR;

        location / {
            proxy_pass http://sd;
            proxy_set_header Upgrade           \$http_upgrade;
            proxy_set_header Connection        \$connection_upgrade;
            proxy_set_header X-Real-IP         \$remote_addr;
            proxy_set_header X-Forwarded-For   \$proxy_add_x_forwarded_for;
        }
    }
}
EOL

# 生成启动脚本
cat <<EOL > run_sd.sh
#!/bin/bash

# 激活虚拟环境
source $venv_name/bin/activate
export HSA_OVERRIDE_GFX_VERSION=$found_version

EOL

# 添加每个实例的启动命令
for i in $(seq 0 $((GPU_COUNT - 1))); do
    port=$((7860 + i))
    cat <<EOL >> run_sd.sh

# 启动第$((i+1))个实例
export HIP_VISIBLE_DEVICES=$i
python3 launch.py --port $port --precision full --no-half --autolaunch --listen --theme dark --skip-torch-cuda-test --server-name $IP_ADDR &
EOL
done

# 确认所有实例已启动
cat <<EOL >> run_sd.sh

echo "所有实例已启动。"
EOL

# 赋予启动脚本执行权限
chmod +x run_sd.sh

echo "启动脚本 run_sd.sh 已生成。"

# 启动并配置Nginx服务为开机自启动
sudo systemctl restart nginx
sudo systemctl enable nginx

echo "Nginx已启动并配置为开机自启动。"
